import jax
import jax.numpy as jnp
from robotisgp import Robotisgp
from learner import *
import defaults
import flax
import functools

@flax.struct.dataclass
class CollectState:
    key: jax.random.PRNGKey
    losses: jnp.ndarray
    n: jnp.ndarray

def run_collecting(collect_state, state, step_fn):
    # pmap要求输入必须是jnp.ndarray，且第一维为显卡设备个数
    key, key_rand = jax.random.split(collect_state.key)

    act = jax.random.uniform(key_rand,(core_env.batch_size, core_env.action_size))
    new_state = step_fn(state, act)
    losses = collect_state.losses
    p_loss = jnp.mean(jnp.mean(jnp.square(state.obs-new_state.obs), axis=-1))

    collect_state = collect_state.replace(key=key, losses=losses + p_loss, n=collect_state.n+1)
    return collect_state, new_state

if __name__ == '__main__':
    local_devices_to_use = jax.local_device_count()
    process_count = jax.process_count()
    env_fn = create_fn(Robotisgp)
    core_env = env_fn(
        action_repeat=defaults.action_repeat,
        batch_size=defaults.num_envs // local_devices_to_use // process_count,
        episode_length=defaults.episode_length)
    key = jax.random.PRNGKey(0)
    key, key_env = jax.random.split(key, 2)
    key_envs = jax.random.split(key_env, local_devices_to_use)
    collect_state = CollectState(key=jnp.stack(jax.random.split(key, local_devices_to_use)), 
                         losses=jnp.zeros((local_devices_to_use,)),
                         n=jnp.zeros((local_devices_to_use,)))
    run_collecting_step = functools.partial(run_collecting, step_fn=core_env.step)
    run_collecting_step = jax.pmap(run_collecting_step, axis_name='i')
    state = jax.tree_multimap(lambda *args: jnp.stack(args),
                                    *[core_env.reset(key) for key in key_envs])
    for i in range(1000):
        collect_state, state = run_collecting_step(collect_state, state)
    print(jnp.mean(collect_state.losses)/collect_state.n)
